import cv2
import numpy as np
import os
from typing import Tuple, Dict, Any, List


class WireDetector:
    """
    铁柱上铁丝缠绕检测器
    使用传统计算机视觉方法检测铁柱上是否有铁丝缠绕
    """

    def __init__(self, config: Dict[str, Any] = None):
        """
        初始化检测器

        Args:
            config: 配置参数字典，包含各种检测参数
        """
        # 默认配置参数
        default_config = {
            # 图像预处理参数
            'blur_kernel_size': 5,  # 高斯模糊核大小
            'bilateral_d': 9,  # 双边滤波直径
            'bilateral_sigma_color': 75,  # 双边滤波颜色空间标准差
            'bilateral_sigma_space': 75,  # 双边滤波坐标空间标准差

            # 边缘检测参数
            'canny_low_threshold': 50,  # Canny边缘检测低阈值
            'canny_high_threshold': 150,  # Canny边缘检测高阈值

            # 霍夫线变换参数
            'hough_rho': 1,  # 距离精度
            'hough_theta': np.pi / 180,  # 角度精度
            'hough_threshold': 100,  # 累加器阈值
            'hough_min_line_length': 50,  # 最小线长
            'hough_max_line_gap': 10,  # 最大线间隙

            # 轮廓分析参数
            'min_contour_area': 100,  # 最小轮廓面积
            'max_contour_area': 50000,  # 最大轮廓面积
            'aspect_ratio_threshold': 0.1,  # 长宽比阈值

            # 铁丝特征参数
            'min_wire_lines': 3,  # 最小铁丝线条数
            'spiral_angle_threshold': 20,  # 螺旋角度阈值（度）
            'line_density_threshold': 0.002  # 线条密度阈值
        }

        # 合并用户配置和默认配置
        self.config = default_config
        if config:
            self.config.update(config)

    def preprocess_image(self, image: np.ndarray) -> np.ndarray:
        """
        图像预处理

        Args:
            image: 输入图像

        Returns:
            预处理后的灰度图像
        """
        # 转换为灰度图
        if len(image.shape) == 3:
            gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
        else:
            gray = image.copy()

        # 双边滤波去噪，保持边缘
        filtered = cv2.bilateralFilter(
            gray,
            self.config['bilateral_d'],
            self.config['bilateral_sigma_color'],
            self.config['bilateral_sigma_space']
        )

        # 高斯模糊进一步平滑
        blurred = cv2.GaussianBlur(
            filtered,
            (self.config['blur_kernel_size'], self.config['blur_kernel_size']),
            0
        )

        return blurred

    def detect_edges(self, image: np.ndarray) -> np.ndarray:
        """
        边缘检测

        Args:
            image: 预处理后的图像

        Returns:
            边缘检测结果
        """
        edges = cv2.Canny(
            image,
            self.config['canny_low_threshold'],
            self.config['canny_high_threshold']
        )
        return edges

    def detect_lines(self, edges: np.ndarray) -> List[np.ndarray]:
        """
        线条检测

        Args:
            edges: 边缘检测结果

        Returns:
            检测到的线条列表
        """
        lines = cv2.HoughLinesP(
            edges,
            self.config['hough_rho'],
            self.config['hough_theta'],
            self.config['hough_threshold'],
            minLineLength=self.config['hough_min_line_length'],
            maxLineGap=self.config['hough_max_line_gap']
        )

        if lines is None:
            return []

        return lines

    def analyze_contours(self, edges: np.ndarray) -> List[Dict]:
        """
        轮廓分析

        Args:
            edges: 边缘检测结果

        Returns:
            轮廓分析结果列表
        """
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        contour_features = []
        for contour in contours:
            area = cv2.contourArea(contour)

            # 过滤太小或太大的轮廓
            if (area < self.config['min_contour_area'] or
                    area > self.config['max_contour_area']):
                continue

            # 计算轮廓特征
            x, y, w, h = cv2.boundingRect(contour)
            aspect_ratio = float(w) / h

            # 计算轮廓的周长
            perimeter = cv2.arcLength(contour, True)

            # 计算紧凑性（圆形度）
            if perimeter > 0:
                compactness = 4 * np.pi * area / (perimeter * perimeter)
            else:
                compactness = 0

            contour_features.append({
                'contour': contour,
                'area': area,
                'aspect_ratio': aspect_ratio,
                'compactness': compactness,
                'bbox': (x, y, w, h)
            })

        return contour_features

    def calculate_line_angles(self, lines: List[np.ndarray]) -> List[float]:
        """
        计算线条角度

        Args:
            lines: 检测到的线条

        Returns:
            线条角度列表（度）
        """
        angles = []
        for line in lines:
            x1, y1, x2, y2 = line[0]
            angle = np.arctan2(y2 - y1, x2 - x1) * 180 / np.pi
            angles.append(abs(angle))

        return angles

    def detect_spiral_pattern(self, lines: List[np.ndarray]) -> bool:
        """
        检测螺旋模式（铁丝缠绕特征）

        Args:
            lines: 检测到的线条

        Returns:
            是否检测到螺旋模式
        """
        if len(lines) < self.config['min_wire_lines']:
            return False

        # 计算线条角度
        angles = self.calculate_line_angles(lines)

        # 检查角度分布是否符合螺旋特征
        angle_variation = np.std(angles)

        # 螺旋线应该有一定的角度变化
        if angle_variation < self.config['spiral_angle_threshold']:
            return False

        return True

    def calculate_line_density(self, lines: List[np.ndarray], image_shape: Tuple[int, int]) -> float:
        """
        计算线条密度

        Args:
            lines: 检测到的线条
            image_shape: 图像形状

        Returns:
            线条密度
        """
        #if not lines:
        #   return 0

        total_length = 0
        for line in lines:
            x1, y1, x2, y2 = line[0]
            length = np.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
            total_length += length

        image_area = image_shape[0] * image_shape[1]
        density = total_length / image_area

        return density

    def detect_wire(self, image_path: str) -> Dict[str, Any]:
        """
        检测铁柱上是否有铁丝缠绕

        Args:
            image_path: 图像文件路径

        Returns:
            检测结果字典，包含是否检测到铁丝和详细信息
        """
        # 读取图像
        if not os.path.exists(image_path):
            return {
                'has_wire': False,
                'error': f'图像文件不存在: {image_path}'
            }

        image = cv2.imread(image_path)
        if image is None:
            return {
                'has_wire': False,
                'error': f'无法读取图像文件: {image_path}'
            }

        # 预处理
        processed_image = self.preprocess_image(image)

        # 边缘检测
        edges = self.detect_edges(processed_image)

        # 线条检测
        lines = self.detect_lines(edges)

        # 轮廓分析
        contour_features = self.analyze_contours(edges)

        # 计算特征
        line_density = self.calculate_line_density(lines, image.shape[:2])
        has_spiral = self.detect_spiral_pattern(lines)

        # 综合判断是否有铁丝缠绕
        has_wire = False
        confidence_score = 0

        # 基于多个特征的综合判断
        if len(lines) >= self.config['min_wire_lines']:
            confidence_score += 0.3

        if line_density > self.config['line_density_threshold']:
            confidence_score += 0.3

        if has_spiral:
            confidence_score += 0.4

        # 基于轮廓特征的补充判断
        thin_contours = [c for c in contour_features
                         if c['aspect_ratio'] < self.config['aspect_ratio_threshold'] or
                         c['aspect_ratio'] > 1 / self.config['aspect_ratio_threshold']]

        if len(thin_contours) > 0:
            confidence_score += 0.2

        # 最终判断
        has_wire = confidence_score >= 0.5

        return {
            'has_wire': has_wire,
            'confidence_score': confidence_score,
            'line_count': len(lines),
            'line_density': line_density,
            'has_spiral_pattern': has_spiral,
            'contour_count': len(contour_features),
            'thin_contour_count': len(thin_contours)
        }

    def batch_detect(self, image_paths: List[str]) -> Dict[str, Dict[str, Any]]:
        """
        批量检测多张图像

        Args:
            image_paths: 图像文件路径列表

        Returns:
            批量检测结果字典
        """
        results = {}
        for image_path in image_paths:
            results[image_path] = self.detect_wire(image_path)

        return results

    def update_config(self, new_config: Dict[str, Any]):
        """
        更新配置参数

        Args:
            new_config: 新的配置参数字典
        """
        self.config.update(new_config)
